from BCEmbedding.tools.langchain import BCERerank
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import UnstructuredPDFLoader

from langchain_community.vectorstores.faiss import FAISS

from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores.utils import DistanceStrategy
from langchain.retrievers import ContextualCompressionRetriever

from pprint import pprint

embedding_model_name = "D:/evns/models/bce/maidalun/bce-embedding-base_v1"
embedding_model_kwargs = {"device":"cuda:0"}
embedding_encode_kwargs = {"batch_size":32,
                           "normalize_embeddings":True,
                           "show_progress_bar":True}

embed_model = HuggingFaceBgeEmbeddings(
    model_name = embedding_model_name,
  #  model_kwargs = embedding_model_kwargs,
    encode_kwargs = embedding_encode_kwargs
)

reranker_args = {
    "model":"D:/evns/models/bce/maidalun/bce-reranker-base_v1",
    "top_n":5,
    "device":"cpu"
}

reranker = BCERerank(**reranker_args)

documents = UnstructuredPDFLoader("files/test.pdf").load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)
texts = text_splitter.split_documents(documents)

print(len(texts))

pprint(texts[0])